import pandas as pd
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from kam import readfile


df = readfile()
# 获取矩阵行数
line = df.shape[0]
# 获取矩阵列数
col = df.shape[1]
data = []
# 行数
for i in range(df.shape[0]):
    if i == 0:
        continue
    else:
        data.append(df.iloc[i])
featureList = df.columns[1: col]
mdl = pd.DataFrame.from_records(data, columns=featureList)
# 存放每次结果的误差平方和
SSE = []
for k in range(1, col):
    # 构造聚类器
    estimator = KMeans(n_clusters=k)
    estimator.fit(
        np.array(mdl[featureList]))
    SSE.append(estimator.inertia_)
# X取决于元素个数
X = range(1, col)
# x轴
plt.xlabel('k')
# y轴
plt.ylabel('SSE')
# '利用SSE选择k'
plt.plot(X, SSE, 'o-')
plt.show()
